package calibration

import (
	"data-treating/pkg/types"
	"math"
	"sync"
	"time"
)

// Calibrator 校准器
type Calibrator struct {
	calibrations    map[string]*types.CalibrationData // 按人员ID存储校准数据
	currentPersonID string                            // 当前活动的人员ID
	minFrames       int
	maxFrames       int
	mutex           sync.RWMutex
}

// NewCalibrator 创建新的校准器
func NewCalibrator() *Calibrator {
	return &Calibrator{
		calibrations: make(map[string]*types.CalibrationData),
		minFrames:    10, // 最少需要10帧进行校准
		maxFrames:    50, // 最多使用50帧进行校准
	}
}

// SetCurrentPerson 设置当前活动的人员ID
func (c *Calibrator) SetCurrentPerson(personID string) {
	c.mutex.Lock()
	defer c.mutex.Unlock()
	c.currentPersonID = personID

	// 如果该人员还没有校准数据，创建一个新的
	if _, exists := c.calibrations[personID]; !exists {
		c.calibrations[personID] = &types.CalibrationData{
			PersonID:       personID,
			ReferencePoses: make([]types.PoseData, 0),
			Calibrated:     false,
			CreatedAt:      time.Now().Unix(),
			UpdatedAt:      time.Now().Unix(),
		}
	}
}

// GetCurrentPersonID 获取当前活动的人员ID
func (c *Calibrator) GetCurrentPersonID() string {
	c.mutex.RLock()
	defer c.mutex.RUnlock()
	return c.currentPersonID
}

// GetCalibrationData 获取当前人员的校准数据
func (c *Calibrator) GetCalibrationData() *types.CalibrationData {
	c.mutex.RLock()
	defer c.mutex.RUnlock()

	if c.currentPersonID == "" {
		return nil
	}

	if data, exists := c.calibrations[c.currentPersonID]; exists {
		return data
	}
	return nil
}

// GetCalibrationDataByPerson 根据人员ID获取校准数据
func (c *Calibrator) GetCalibrationDataByPerson(personID string) *types.CalibrationData {
	c.mutex.RLock()
	defer c.mutex.RUnlock()

	if data, exists := c.calibrations[personID]; exists {
		return data
	}
	return nil
}

// SetCalibrationData 设置指定人员的校准数据
func (c *Calibrator) SetCalibrationData(personID string, data *types.CalibrationData) {
	c.mutex.Lock()
	defer c.mutex.Unlock()

	if data != nil {
		data.PersonID = personID
		data.UpdatedAt = time.Now().Unix()
		c.calibrations[personID] = data
	}
}

// GetAllCalibrations 获取所有人员的校准数据
func (c *Calibrator) GetAllCalibrations() map[string]*types.CalibrationData {
	c.mutex.RLock()
	defer c.mutex.RUnlock()

	result := make(map[string]*types.CalibrationData)
	for personID, data := range c.calibrations {
		result[personID] = data
	}
	return result
}

// DeleteCalibration 删除指定人员的校准数据
func (c *Calibrator) DeleteCalibration(personID string) bool {
	c.mutex.Lock()
	defer c.mutex.Unlock()

	if _, exists := c.calibrations[personID]; exists {
		delete(c.calibrations, personID)
		return true
	}
	return false
}

// AddCalibrationFrame 添加校准帧
func (c *Calibrator) AddCalibrationFrame(pose types.PoseData) bool {
	c.mutex.Lock()
	defer c.mutex.Unlock()

	if c.currentPersonID == "" {
		return false
	}

	calibrationData := c.calibrations[c.currentPersonID]
	if calibrationData == nil {
		return false
	}

	// 检查关键点是否有效
	if !c.isPoseValid(pose) {
		return false
	}

	// 限制校准帧数量
	if len(calibrationData.ReferencePoses) >= c.maxFrames {
		// 移除最旧的帧
		calibrationData.ReferencePoses = calibrationData.ReferencePoses[1:]
	}

	calibrationData.ReferencePoses = append(calibrationData.ReferencePoses, pose)
	calibrationData.UpdatedAt = time.Now().Unix()
	return true
}

// IsCalibrationReady 检查是否可以进行校准
func (c *Calibrator) IsCalibrationReady() bool {
	c.mutex.RLock()
	defer c.mutex.RUnlock()

	if c.currentPersonID == "" {
		return false
	}

	calibrationData := c.calibrations[c.currentPersonID]
	if calibrationData == nil {
		return false
	}

	return len(calibrationData.ReferencePoses) >= c.minFrames
}

// PerformCalibration 执行校准
func (c *Calibrator) PerformCalibration() error {
	c.mutex.Lock()
	defer c.mutex.Unlock()

	if c.currentPersonID == "" {
		return ErrInsufficientFrames
	}

	calibrationData := c.calibrations[c.currentPersonID]
	if calibrationData == nil {
		return ErrInsufficientFrames
	}

	if len(calibrationData.ReferencePoses) < c.minFrames {
		return ErrInsufficientFrames
	}

	// 计算平均姿势
	c.calculateAveragePose(calibrationData)

	// 计算标准差
	c.calculateStandardDeviations(calibrationData)

	calibrationData.Calibrated = true
	calibrationData.UpdatedAt = time.Now().Unix()
	return nil
}

// ResetCalibration 重置校准数据
func (c *Calibrator) ResetCalibration() {
	c.mutex.Lock()
	defer c.mutex.Unlock()

	if c.currentPersonID == "" {
		return
	}

	c.calibrations[c.currentPersonID] = &types.CalibrationData{
		PersonID:       c.currentPersonID,
		ReferencePoses: make([]types.PoseData, 0),
		Calibrated:     false,
		CreatedAt:      time.Now().Unix(),
		UpdatedAt:      time.Now().Unix(),
	}
}

// calculateAveragePose 计算平均姿势
func (c *Calibrator) calculateAveragePose(calibrationData *types.CalibrationData) {
	if len(calibrationData.ReferencePoses) == 0 {
		return
	}

	var avgPose types.PoseData
	avgPose.Timestamp = time.Now().UnixNano()
	avgPose.FrameID = 0

	// 计算每个关键点的平均值
	for i := 0; i < 17; i++ {
		var sumX, sumY, sumConf float64
		validCount := 0

		for _, pose := range calibrationData.ReferencePoses {
			if pose.Keypoints[i].IsValid() {
				sumX += pose.Keypoints[i].X
				sumY += pose.Keypoints[i].Y
				sumConf += pose.Keypoints[i].Confidence
				validCount++
			}
		}

		if validCount > 0 {
			avgPose.Keypoints[i].X = sumX / float64(validCount)
			avgPose.Keypoints[i].Y = sumY / float64(validCount)
			avgPose.Keypoints[i].Confidence = sumConf / float64(validCount)
		}
	}

	calibrationData.AveragePose = avgPose
}

// calculateStandardDeviations 计算标准差
func (c *Calibrator) calculateStandardDeviations(calibrationData *types.CalibrationData) {
	if len(calibrationData.ReferencePoses) == 0 {
		return
	}

	// 计算每个关键点的标准差
	for i := 0; i < 17; i++ {
		var sumX, sumY, sumXSq, sumYSq float64
		validCount := 0

		avgX := calibrationData.AveragePose.Keypoints[i].X
		avgY := calibrationData.AveragePose.Keypoints[i].Y

		for _, pose := range calibrationData.ReferencePoses {
			if pose.Keypoints[i].IsValid() {
				dx := pose.Keypoints[i].X - avgX
				dy := pose.Keypoints[i].Y - avgY
				sumX += dx
				sumY += dy
				sumXSq += dx * dx
				sumYSq += dy * dy
				validCount++
			}
		}

		if validCount > 1 {
			// 计算标准差
			stdX := math.Sqrt((sumXSq - sumX*sumX/float64(validCount)) / float64(validCount-1))
			stdY := math.Sqrt((sumYSq - sumY*sumY/float64(validCount)) / float64(validCount-1))

			calibrationData.StandardDeviations[i].X = stdX
			calibrationData.StandardDeviations[i].Y = stdY
			calibrationData.StandardDeviations[i].Confidence = 1.0
		}
	}
}

// isPoseValid 检查姿势数据是否有效
func (c *Calibrator) isPoseValid(pose types.PoseData) bool {
	validKeypoints := 0
	minValidKeypoints := 12 // 至少需要12个有效关键点

	for _, keypoint := range pose.Keypoints {
		if keypoint.IsValid() {
			validKeypoints++
		}
	}

	return validKeypoints >= minValidKeypoints
}

// GetCalibrationProgress 获取校准进度
func (c *Calibrator) GetCalibrationProgress() float64 {
	c.mutex.RLock()
	defer c.mutex.RUnlock()

	if c.currentPersonID == "" {
		return 0
	}

	calibrationData := c.calibrations[c.currentPersonID]
	if calibrationData == nil {
		return 0
	}

	if c.minFrames == 0 {
		return 0
	}
	progress := float64(len(calibrationData.ReferencePoses)) / float64(c.minFrames)
	return math.Min(progress, 1.0)
}
